#include <stdio.h>
#include <stdlib.h>
#include "RBTree.h"

/* 将x进行左旋，将左、右、父节点进行更新
 *    px                              px
 *    |                               |
 *    x                                y
 *   /  \      --(左旋)-->    / \
 *  lx   y                          x  ry
 *     /   \                     /  \
 *    ly   ry                   lx  ly
*/
static void leftRotate(RBTree *tree, RBNode *x) {
	RBNode *y = x->right;
	x->right = y->left;
	if (y->left) {
		y->left->parent = x;
	}        // 将x与ly接上
	y->parent = x->parent;  // 将x的位置异位给y
	if (x->parent) {
		if (x->parent->left == x) {
			x->parent->left = y;  
		} else {
			x->parent->right = y;   // px和y接上
		}
	} else {
		tree->root = y;  // 否则顶点为y
	}
	y->left = x;
	x->parent = y;
}

/* 将y进行左旋，将左、右、父节点进行更新
 *          py                               py
 *          |                                |
 *          y                                x
 *         /  \      --(右旋)-->            /  \
 *        x   ry                           lx   y
 *       / \                                   / \
 *      lx  rx                                rx  ry
 * */
static void rightRotate(RBTree* tree, RBNode *y) {
	RBNode *x = y->left;
	y->left = x->right;
	if (x->right) {
		x->right->parent = y;
	}
	x->parent = y->parent;
	if (y->parent) {
		if (y->parent->right == y) {
			y->parent->right = x;
		} else {
			y->parent->left =x;
		}
	} else {
		tree->root = x;
	}
	x->right = y;
	y->parent = x;
}

static RBNode *createRBNode(KeyType key) {
	RBNode *node = malloc(sizeof(RBNode));
	node->color = RED;
	node->key = key;
	node->left = node->right = node->parent = NULL;
	return node;
}

static void destoryRBNode(RBNode *node) {
	if (node) {
		destoryRBNode(node->left);
		destoryRBNode(node->right);
		free(node);
	}
}
/* 1. 插入元素，如果父节点是黑色，那么不用调整
 * 2. 如果父节点是红色，有连续的两个红
 * 2.1 叔叔节点是红色
 *		重新着色(p->黑色, u->黑色， g->红色)，作为新节点重新进行判断
 * 2.2 叔叔节点是黑色
 *		2.2.1 cur左孩子，par左孩子
 *			g右旋，g->红色，p->黑色
 *		2.2.2 cur右孩子，par右孩子
 *			g左旋，g->红色，p->黑色
 *		2.2.3 cur右孩子，par左孩子
 *			p左旋，cur和par交换，然后重复2.2.1
 *		2.2.4 cur左孩子，par右孩子
 *			p右旋，cur和par交换，然后重复2.2.2
 */
static void insertFixUp(RBTree *tree, RBNode *node) {
	RBNode *parent = NULL, *grandParent = NULL;
	RBNode *uncle = NULL;
	RBNode *tmp;
	parent = node->parent;
	// 判断父节点是否为红色，父节点是黑色直接插入不需要调整
	while (parent && parent->color == RED) {
		// 找祖父节点  -- 为了找叔叔
		grandParent = parent->parent;
		if (parent == grandParent->left) {
			uncle = grandParent->right;
		} else {
			uncle = grandParent->left;
		}

		if (uncle && uncle->color == RED) {   // 叔叔是红色
			uncle->color = BLACK;
			parent->color = BLACK;
			grandParent->color = RED;	// 叔父爷变色 + cur换指向
			node = grandParent;
			parent = node->parent;  // 将父节点指向当前节点的父节点，因为当前结点变成了grandParent
			continue;
		}
		// 叔叔是黑色
		if (grandParent->left == parent) {		// L，左边为父节点
			if (parent->right == node) {		// R
				leftRotate(tree, parent);		// 左旋父节点
				tmp = parent;				// 交换父节点和插入的 node 节点
				parent = node;
				node = tmp;
			}
			// LL
			rightRotate(tree, grandParent);	// 旋转爷爷，旋转点和旋转中心的进行变色（旋转点是爷爷节点，旋转中心点是父节点）
			grandParent->color = RED;
			parent->color = BLACK;
		} else {							// 同上
			if (parent->left == node) {		// RL
				rightRotate(tree, parent);
				tmp = parent;
				parent = node;
				node = tmp;
			} // RR
			leftRotate(tree, grandParent);
			grandParent->color = RED;
			parent->color = BLACK;
		}

	}
	tree->root->color = BLACK;   // 将根节点置为黑
}

RBTree *initRBTree() {
	RBTree *tree = malloc(sizeof(RBTree));
	tree->count = 0;
	tree->root = NULL;
	return tree;
}

void releaseRBTree(RBTree* tree) {
	if (tree) {
		destoryRBNode(tree->root);
		free(tree);
	}
}

void insertRBTree(RBTree *tree, KeyType key) {
	// 1. 先创建一个红色节点
	RBNode *node = createRBNode(key);
	// 2. 根据二叉搜索树的规则找到待插入的位置
	RBNode *cur = tree->root;  // 大哥
	RBNode *pre = NULL;        // 小弟
	while (cur) {
		pre = cur;
		if (key < cur->key) {
			cur = cur->left;
		} else if (key > cur->key) {
			cur = cur->right;
		} else {
			printf("Key: %d have exist!\n", key);
		}
	}

	// 3. 在对应的位置上插入，若是根，更新tree
	node->parent = pre;
	if (pre) {
		if (key < pre->key) {
			pre->left = node;
		} else if (key > pre->key) {
			pre->right = node;
		}
	} else {
		tree->root = node;  // 如果没有根
	}
	// 4. 修改红黑树
	insertFixUp(tree, node);
}

static void showRBTree(RBNode *node, int key, int dir) {
	if (node) {
		if (dir == 0) {
			printf("%02d[B] is root\n", node->key);
		} else {
			printf("%2d[%c] is %2d's %s\n", node->key,
				(node->color == RED) ? 'R': 'B', key,
				(dir == 1) ? "right child": "left child");
		}
		showRBTree(node->left, node->key, -1);
		showRBTree(node->right, node->key, 1);
	}
}

void printRBTree(RBTree *tree) {
	if (tree) {
		showRBTree(tree->root, tree->root->key, 0);
	}
}

RBNode* searchRBNode(RBTree* tree, KeyType key) {
	RBNode* node = tree->root;
	while (node) {
		if (key < node->key) {
			node = node->left;
		} else if (key > node->key) {
			node = node->right;
		} else {
			return node;
		}
	}
	return NULL;
}

/* 调整的方法：
1. 兄弟节点 是黑色的
	1.1 兄弟的同向节点有红色
	1.2 兄弟的异向节点有红色
	1.3 兄弟节点下的两个节点都是黑色
		把兄弟节点变为红色，双黑节点的多余的黑色属性 传递给 父节点
		父节点变为双黑，递归
2. 兄弟节点 是红色的
	 兄弟节点变为黑色
 */
static void deleteFixup(RBTree *tree, RBNode *x, RBNode *parent) {
	RBNode *w;   // 兄弟结点
	while ((tree->root != x) && (x == NULL || x->color == BLACK)) {
		if (parent->left == x) {
			w = parent->right;          // 兄弟在右边
			if (w->color == RED) {      // 兄弟是红色
				w->color = BLACK;		// (1) x的兄弟节点设置为黑色
				parent->color = RED;	// (2) x的父节点设置为红色
				leftRotate(tree, parent);// (3) 对x节点的父节点进行左旋
				w = parent->right;
			}
			// 兄弟的子节点都是黑色的
			if ((w->left == NULL || w->left->color == BLACK) &&
				(w->right == NULL || w->right->color == BLACK)) {   // 没有孩子或兄弟孩子都是黑色
				w->color = RED;             // 删除的结点为为当前结点，如果删除之后，则将删除节点的父节点置为当前结点
				x = parent;
				parent = x->parent;        // 祖父节点变为父节点了
			} else {  // 兄弟至少一个红孩子
				if (w->right == NULL || w->right->color == BLACK) {
					w->left->color = BLACK;
					w->color = RED;
					rightRotate(tree, w);
					w = parent->right;  // 相当于RL中的R
				}
				w->color = parent->color;
				parent->color = BLACK;
				w->right->color = BLACK;
				leftRotate(tree, parent);
				x = tree->root;    // 相当于LL
				break;
			}
		} else {   // 兄弟在左边
			w = parent->left;
			if (w->color == RED) { 
				w->color = BLACK;
				parent->color = RED;
				rightRotate(tree, parent);
				w = parent->left;
			}
			if ((w->left == NULL || w->left->color == BLACK) &&
				(w->right == NULL || w->right->color == BLACK)) {
				w->color = RED;
				x = parent;
				parent = x->parent;
			} else {
				if (w->left == NULL || w->left->color == BLACK) {
					w->right->color = BLACK;
					w->color = RED;
					leftRotate(tree, w);
					w = parent->left;
				}
				w->color = parent->color;
				parent->color = BLACK;
				w->left->color = BLACK;
				rightRotate(tree, parent);
				x = tree->root;
				break;
			}
		}
	}
	if (x) {
		x->color = BLACK;
	}
}

static void deleteRBNode(RBTree *tree, RBNode *node) {
	RBNode *y;				// 真正删除的节点地址
	RBNode *x;				// 替换节点
	RBNode *parent = NULL;	// 替换节点是NULL，无法找到父节点

	if (node->left == NULL || node->right == NULL) {
		y = node;
	} else {  // 寻找的是后继
		y = node->right;
		while (y->left) {
			y = y->left;
		}
	}
	// 真正删除的节点找到了，开始寻找替换节点
	if (y->left) {
		x = y->left;
	} else {
		x = y->right;
	}
	parent = y->parent;
	// 开始更新替换节点和父节点的关系
	if (x) {
		x->parent = y->parent;
	}
	if (y->parent == NULL) {
		tree->root = x;
	} else if (y->parent->left == y) {
		y->parent->left = x;
	} else {
		y->parent->right = x;
	}
	// 如果度为2，将待删除的值赋值给度为2的点
	if (node != y) {
		node->key = y->key;
	}
	if (y->color == BLACK) {
		deleteFixup(tree, x, parent);
	}
	free(y);
}

void deleteRBTree(RBTree *tree, KeyType key) {
	// 1. 查找到待删除的节点
	RBNode *node = searchRBNode(tree, key);
	// 2. 处理x和y的关系
	if (node) {
		deleteRBNode(tree, node);
	}
}
